#include "c_ldtw.h"
#include <math.h>
#include <stdlib.h>
#include <limits.h>

/* 内部工具：max_steps */
static int max_steps(int i, int j, int max_len, int len1, int len2)
{
    int c1 = i + j;
    int c2 = max_len - ((len1 - i - 1) > (len2 - j - 1) ? (len1 - i - 1) : (len2 - j - 1));
    return c1 < c2 ? c1 : c2;
}

/* 主算法：带长度限制的 DTW；返回结果需外部释放 */
LDTWResult ldtw_path(const double *s1, size_t len1,
                     const double *s2, size_t len2,
                     int max_length)
{
    /* 参数检查 */
    if (max_length < (int)(len1 > len2 ? len1 : len2))
    {
        /* 用 path==NULL 表示异常 */
        return (LDTWResult){NULL, 0, -1.0};
    }

    /* 分配三维代价数组：cost[i][j][s] 用扁平索引 */
    size_t stride = (size_t)(max_length + 1);
    size_t mat_size = len1 * len2 * stride;
    double *cost = (double *)calloc(mat_size, sizeof(double));
    if (!cost)
    {
        return (LDTWResult){NULL, 0, -1.0};
    }
    /* 初始化成无穷大 */
    for (size_t k = 0; k < mat_size; ++k)
    {
        cost[k] = INFINITY;
    }
    /* 初始化 (0,0) */
    cost[0] = (s1[0] - s2[0]) * (s1[0] - s2[0]);

    /* 第一列 */
    for (size_t i = 1; i < len1; ++i)
    {
        size_t idx = (i * len2 + 0) * stride + i;
        double prev = cost[((i - 1) * len2 + 0) * stride + (i - 1)];
        cost[idx] = prev + (s1[i] - s2[0]) * (s1[i] - s2[0]);
    }
    /* 第一行 */
    for (size_t j = 1; j < len2; ++j)
    {
        size_t idx = (0 * len2 + j) * stride + j;
        double prev = cost[(0 * len2 + (j - 1)) * stride + (j - 1)];
        cost[idx] = prev + (s1[0] - s2[j]) * (s1[0] - s2[j]);
    }

    /* DP 主体 */
    for (size_t i = 1; i < len1; ++i)
    {
        for (size_t j = 1; j < len2; ++j)
        {
            int min_s = (i > j ? (int)i : (int)j);
            int max_s = max_steps((int)i, (int)j, max_length - 1,
                                  (int)len1, (int)len2);
            double d = (s1[i] - s2[j]) * (s1[i] - s2[j]);

            for (int s = min_s; s <= max_s; ++s)
            {
                size_t idx = (i * len2 + j) * stride + s;

                /* 三个前驱 */
                double diag = cost[((i - 1) * len2 + (j - 1)) * stride + (s - 1)];
                double left = cost[((i - 1) * len2 + j) * stride + (s - 1)];
                double up = cost[(i * len2 + (j - 1)) * stride + (s - 1)];

                double best = diag;
                if (left < best)
                    best = left;
                if (up < best)
                    best = up;

                cost[idx] = d + best;
            }
        }
    }

    /* 在终点 (len1-1, len2-1) 找最优长度 */
    size_t last_i = len1 - 1, last_j = len2 - 1;
    double best_cost = 1e100;
    int best_s = -1;
    for (int s = 0; s <= max_length; ++s)
    {
        double v = cost[(last_i * len2 + last_j) * stride + s];
        if (v < best_cost)
        {
            best_cost = v;
            best_s = s;
        }
    }
    /* 回溯前再检查 best_s 合理性 */
    if (best_s < 0 || best_s > max_length)
    {
        free(cost);
        return (LDTWResult){NULL, 0, -1.0};
    }
    /* 回溯路径（用临时动态数组） */
    PathNode *rev_path = (PathNode *)malloc(sizeof(PathNode) * (best_s + 1));
    int rp = 0;
    int ci = (int)last_i, cj = (int)last_j, cs = best_s;
    while (ci != 0 || cj != 0)
    {
        rev_path[rp++] = (PathNode){ci, cj};
        if (ci == 0)
        {
            cj--;
        }
        else if (cj == 0)
        {
            ci--;
        }
        else
        {
            /* 选最小前驱 */
            double diag = cost[((ci - 1) * len2 + (cj - 1)) * stride + (cs - 1)];
            double left = cost[((ci - 1) * len2 + cj) * stride + (cs - 1)];
            double up = cost[(ci * len2 + (cj - 1)) * stride + (cs - 1)];

            if (diag <= left && diag <= up)
            {
                ci--;
                cj--;
                cs--;
            }
            else if (left <= diag && left <= up)
            {
                ci--;
                cs--;
            }
            else
            {
                cj--;
                cs--;
            }
        }
    }
    rev_path[rp++] = (PathNode){0, 0};

    /* 反转到正序 */
    PathNode *path = (PathNode *)malloc(sizeof(PathNode) * rp);
    for (int k = 0; k < rp; ++k)
    {
        path[k] = rev_path[rp - 1 - k];
    }
    free(rev_path);

    free(cost);

    LDTWResult ret = {path, rp, sqrt(best_cost)};
    return ret;
}
/* ReLU 惩罚 */
static inline double penalty_s(double k, double beta, double gap_max)
{
    double relu = (k > gap_max) ? (k - gap_max) : 0.0;
    return beta * relu * relu;
}
/* 分段时间加权 LDTW：ReLU 惩罚 */
LDTWResult ldtw_path_tw(const double *s1, size_t len1,
                        const double *s2, size_t len2,
                        int max_length,
                        double beta,
                        const double *gap_max_vec)
{
    if (max_length < (int)(len1 > len2 ? len1 : len2))
        return (LDTWResult){NULL, 0, -1.0};

    size_t stride = (size_t)(max_length + 1);
    size_t mat_size = len1 * len2 * stride;
    double *cost = (double *)calloc(mat_size, sizeof(double));
    if (!cost)
        return (LDTWResult){NULL, 0, -1.0};
    for (size_t k = 0; k < mat_size; ++k)
        cost[k] = INFINITY;

    cost[0] = (s1[0] - s2[0]) * (s1[0] - s2[0]);
    for (size_t i = 1; i < len1; ++i)
        cost[(i * len2 + 0) * stride + i] =
            cost[((i - 1) * len2 + 0) * stride + (i - 1)] +
            (s1[i] - s2[0]) * (s1[i] - s2[0]);
    for (size_t j = 1; j < len2; ++j)
        cost[(0 * len2 + j) * stride + j] =
            cost[(0 * len2 + (j - 1)) * stride + (j - 1)] +
            (s1[0] - s2[j]) * (s1[0] - s2[j]);

    for (size_t i = 1; i < len1; ++i)
    {
        for (size_t j = 1; j < len2; ++j)
        {
            int min_s = (i > j ? (int)i : (int)j);
            int max_s = max_steps((int)i, (int)j, max_length - 1,
                                  (int)len1, (int)len2);
            double geo_d = (s1[i] - s2[j]) * (s1[i] - s2[j]);

            for (int s = min_s; s <= max_s; ++s)
            {
                size_t idx = (i * len2 + j) * stride + s;
                double diag = cost[((i - 1) * len2 + (j - 1)) * stride + (s - 1)];
                double left = cost[((i - 1) * len2 + j) * stride + (s - 1)];
                double up = cost[(i * len2 + (j - 1)) * stride + (s - 1)];
                double best = diag;
                if (left < best)
                    best = left;
                if (up < best)
                    best = up;
                cost[idx] = geo_d + best; /* 只用几何距离 */
            }
        }
    }

    size_t last_i = len1 - 1, last_j = len2 - 1;
    int best_s = -1;
    double best_geo = INFINITY;
    for (int s = 0; s <= max_length; ++s)
    {
        double v = cost[(last_i * len2 + last_j) * stride + s];
        if (v < best_geo)
        {
            best_geo = v;
            best_s = s;
        }
    }
    if (best_s < 0 || best_s > max_length)
    {
        free(cost);
        return (LDTWResult){NULL, 0, -1.0};
    }

    PathNode *rev_path = (PathNode *)malloc(sizeof(PathNode) * (best_s + 1));
    int rp = 0;
    int ci = (int)last_i, cj = (int)last_j, cs = best_s;
    while (ci != 0 || cj != 0)
    {
        rev_path[rp++] = (PathNode){ci, cj};
        if (ci == 0)
            cj--;
        else if (cj == 0)
            ci--;
        else
        {
            double diag = cost[((ci - 1) * len2 + (cj - 1)) * stride + (cs - 1)];
            double left = cost[((ci - 1) * len2 + cj) * stride + (cs - 1)];
            double up = cost[(ci * len2 + (cj - 1)) * stride + (cs - 1)];
            if (diag <= left && diag <= up)
            {
                ci--;
                cj--;
                cs--;
            }
            else if (left <= diag && left <= up)
            {
                ci--;
                cs--;
            }
            else
            {
                cj--;
                cs--;
            }
        }
    }
    rev_path[rp++] = (PathNode){0, 0};
    PathNode *path = (PathNode *)malloc(sizeof(PathNode) * rp);
    for (int k = 0; k < rp; ++k)
        path[k] = rev_path[rp - 1 - k];
    free(rev_path);

    /* 仅对路径中实际产生时间拉伸的阶段施加惩罚 */
    double penalty_sum = 0.0;
    for (int k = 0; k < rp - 1; ++k)
    {
        int di = path[k + 1].i - path[k].i;
        int dj = path[k + 1].j - path[k].j;
        if (di != 0 && dj != 0)
            continue; /* 对角步，无拉伸 */
        /* 单边步，出现拉伸 */
        double gap = fabs((double)path[k].i - (double)path[k].j);
        penalty_sum += penalty_s(gap, beta, gap_max_vec[path[k].i]);
    }
    free(cost);

    LDTWResult ret = {path, rp, sqrt(best_geo + penalty_sum)};
    return ret;
}